#ifndef COARSEN_EDGE_HPP
#define COARSEN_EDGE_HPP
#include "utils.hpp"
#include "coarsen.hpp"
#include <algorithm>
#include <climits>

template<typename idx_t, typename data_t>
struct COO {
    idx_t row_id = -1;
    idx_t col_id = -1;
    data_t v = 0.0;
    COO() {}
    COO(idx_t _r, idx_t _c, data_t _v): row_id(_r), col_id(_c), v(_v) {}
    bool operator < (const COO & b) const {
        if (v == b.v) return (row_id == b.row_id) ? (col_id < b.col_id) : (row_id < b.row_id);
        else return v < b.v;
    }
};


template<typename idx_t, typename data_t, typename calc_t> static
CoarsenWrapper<idx_t, data_t, calc_t>* pairwise_coarsen(const MyParCSRMatrix<idx_t, data_t, calc_t> & fine_mat,
    const calc_t alpha)
{
    int my_pid; MPI_Comm_rank(fine_mat.comm, &my_pid);
    int nprocs; MPI_Comm_size(fine_mat.comm, &nprocs);
    if (my_pid == 0) printf("Coord coarsen\n");
    const idx_t num_fine = fine_mat.end_row - fine_mat.beg_row;// 本进程在细网格所负责的行数

    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_diag = *(fine_mat.diag);
    idx_t * loc_map_f2c = new idx_t [num_fine];// 记录本地的第i个细点被映射为粗点的序号
    for (idx_t fi = 0; fi < num_fine; fi++)
        loc_map_f2c[fi] = -1;// 初始化映射

    const Point<data_t> * fpoints = fine_mat.points; assert(fpoints != nullptr);
    Point<data_t> * cpoints = new Point<data_t> [num_fine];// 先按最大可能开辟
    idx_t num_coar = 0;// 粗点数目
#if 1
    // 先按最大可能开辟
    idx_t * rstr_rpt = new idx_t [num_fine + 1]; rstr_rpt[num_coar] = 0;
    idx_t * rstr_cid = new idx_t [num_fine];
    idx_t  * near_ngbs = new idx_t  [fine_diag.nnz];
    data_t * distances = new data_t [fine_diag.nnz];
    for (idx_t fi = 0; fi < num_fine; fi++) {
        const idx_t pbeg = fine_diag.row_ptr[fi],
                    pend = fine_diag.row_ptr[fi+1];
        for (idx_t p = pbeg; p < pend; p++) {
            const idx_t fj = fine_diag.col_idx[p];
            near_ngbs[p] = fj;
            distances[p] = fpoints[fi].dist(fpoints[fj]);
        }
        bubble_sort<data_t, idx_t>(pend - pbeg, distances + pbeg, near_ngbs + pbeg);
        assert(near_ngbs[pbeg] == fi);

        if (loc_map_f2c[fi] >= 0) continue;// 之前已经聚合过这个点了

        const idx_t coar_id  = num_coar;
        idx_t coar_off = rstr_rpt[num_coar];
        loc_map_f2c[fi] = coar_id;
        rstr_cid[coar_off ++] = fi;

        Point another = fpoints[fi];
        // bool aggr = false;// 是否找到可以聚合的邻居
        for (int p = pbeg + 1; p < pend; p++) {// 从最近的开始找
            const int fj = near_ngbs[p];
            if (loc_map_f2c[fj] >= 0) continue;// 该邻居之前已经被聚合过了
            else {
                // aggr = true;// 可以被聚合
                another = fpoints[fj];
                loc_map_f2c[fj] = coar_id;
                rstr_cid[coar_off ++] = fj;
                break;
            }
        }
        rstr_rpt[num_coar + 1] = coar_off;
        cpoints [num_coar] = fpoints[fi];//fpoints[fi].middle(another);
        num_coar ++;// 聚合得到的粗点数增加
    }
    delete [] near_ngbs; near_ngbs = nullptr;
    delete [] distances; distances = nullptr;
#else
    std::vector<COO<idx_t, data_t> > pair_dists;
    for (idx_t fi = 0; fi < num_fine; fi++) {
        const idx_t pbeg = fine_diag.row_ptr[fi],
                    pend = fine_diag.row_ptr[fi+1];
        for (idx_t p = pbeg; p < pend; p++) {
            const idx_t fj = fine_diag.col_idx[p];
            if (fi >= fj) continue;// == 的不必考虑，>的之前已经添加过（对称性）

            pair_dists.emplace_back(fi, fj, fpoints[fi].dist(fpoints[fj]));
        }
    }

    std::sort(pair_dists.begin(), pair_dists.end());

    idx_t * assum_R_rp = new idx_t [num_fine + 1];// 限制矩阵的行偏移，按最大可能空间开辟
    idx_t * assum_R_ci = new idx_t [num_fine];// 限制矩阵的列索引，它的非零元数一定是细网格的数目
    assum_R_rp[0] = 0;
    for (size_t k = 0; k < pair_dists.size(); k++) {
        const idx_t fi = pair_dists[k].row_id,
                    fj = pair_dists[k].col_id;
        if (loc_map_f2c[fi] == -1 && loc_map_f2c[fj] == -1) {// 这两个点之前都还没被聚合过
            const idx_t coar_id  = num_coar; assert(num_coar + 1 <= num_fine);
            idx_t coar_off = assum_R_rp[num_coar]; assert(coar_off + 2 <= num_fine);
            loc_map_f2c[fi] = coar_id; assum_R_ci[coar_off ++] = fi;
            loc_map_f2c[fj] = coar_id; assum_R_ci[coar_off ++] = fj;
            assum_R_rp[num_coar + 1] = coar_off;
            // printf("%d %d\n", fi, fj);
            num_coar ++;
        }
    }
    for (idx_t fi = 0; fi < num_fine; fi++) {// 剩下的每个细点都对应一个粗点
        if (loc_map_f2c[fi] == -1) {
            const idx_t coar_id  = num_coar;
            idx_t coar_off = assum_R_rp[num_coar];
            loc_map_f2c[fi] = coar_id; assum_R_ci[coar_off ++] = fi;
            assum_R_rp[num_coar + 1] = coar_off;
            // printf("%d -1\n", fi);
            num_coar ++;
        }
    }
    assert(assum_R_rp[num_coar] == num_fine);
    assert(num_coar > 0);// 暂时不处理跨进程粗化
    // 将细点编号小的往前挪，使其粗点编号也较小
    struct F2OldC {
        idx_t minF = -1;
        idx_t oldC = -1;
        bool operator < (const F2OldC & b) const {
            assert(minF != b.minF);
            return minF  < b.minF;
        }
    };
    std::vector<F2OldC> aux_arr(num_coar);
    for (idx_t oci = 0; oci < num_coar; oci++) {
        aux_arr[oci].oldC = oci;
        idx_t minF = INT_MAX;
        for (idx_t p = assum_R_rp[oci]; p < assum_R_rp[oci + 1]; p++)
            minF = MIN(minF, assum_R_ci[p]);
        aux_arr[oci].minF = minF;
    }
    std::sort(aux_arr.begin(), aux_arr.end());
    idx_t * rstr_rpt = new idx_t [num_coar + 1];
    idx_t * rstr_cid = new idx_t [num_fine];
    rstr_rpt[0] = 0;
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        const idx_t pbeg = assum_R_rp[oci], pend = assum_R_rp[oci + 1];
        // Point<data_t> cc;// 粗点的坐标
        for (idx_t p = pbeg; p < pend; p++) {
            idx_t fi = assum_R_ci[p];
            loc_map_f2c[fi] = nci;// 重新构造细到粗的映射（相当于插值矩阵）
            // cc.add(fpoints[fi]);
        }
        // 同时修改限制矩阵的行偏移
        rstr_rpt[nci + 1] = rstr_rpt[nci] + pend - pbeg;
        // 记录粗点坐标
        cpoints[nci] = fpoints[assum_R_ci[pbeg]];
        // cpoints[nci] = cc.div((data_t)(pend - pbeg));
    }
    // 修改限制矩阵的列索引
    for (idx_t nci = 0; nci < num_coar; nci++) {
        idx_t oci = aux_arr[nci].oldC;
        const idx_t num = assum_R_rp[oci + 1] - assum_R_rp[oci];
        idx_t * dst = rstr_cid + rstr_rpt[nci];
        const idx_t * src = assum_R_ci + assum_R_rp[oci];
        for (idx_t p = 0; p < num; p++) {
            dst[p] = src[p];
            // printf("%d ", dst[p]);
        }
        // printf("\n");
    }
    delete [] assum_R_rp; assum_R_rp = nullptr;
    delete [] assum_R_ci; assum_R_ci = nullptr;

    // MPI_Barrier(MPI_COMM_WORLD);
    // MPI_Abort(MPI_COMM_WORLD, -303);
#endif
    // 至此已经获得了聚合的列表

    // 计算粗网格矩阵
    MyParCSRMatrix<idx_t, data_t, calc_t> * coar_mat = new MyParCSRMatrix<idx_t, data_t, calc_t>(fine_mat.comm);
    coar_mat->points = cpoints;
    idx_t * & rows_partition = coar_mat->rows_partition;
    rows_partition = new idx_t [nprocs + 1];
    rows_partition[0] = 0;
    constexpr MPI_Datatype mpi_idx_type = (sizeof(idx_t) == 4) ? MPI_INT : MPI_LONG_LONG;
    MPI_Allgather(&num_coar, 1, mpi_idx_type, rows_partition + 1, 1, mpi_idx_type, fine_mat.comm);
    for (int p = 1; p <= nprocs; p++) {
        rows_partition[p] += rows_partition[p - 1];
    }
    coar_mat->beg_row   = rows_partition[my_pid    ];
    coar_mat->end_row   = rows_partition[my_pid + 1]; assert(coar_mat->end_row - coar_mat->beg_row == num_coar);
    coar_mat->glb_nrows = rows_partition[nprocs];
    // 对角和非对角部分的矩阵值可以分别计算
    aggregate_mat<idx_t, data_t, calc_t>(fine_diag, num_coar, num_coar, rstr_rpt, rstr_cid, loc_map_f2c, coar_mat->diag, alpha);
    // record.emplace_back(MPI_Wtime());

    const MyCSRMatrix<idx_t, data_t, calc_t> & fine_offd = *(fine_mat.offd);
    idx_t * remote_map_f2c = new idx_t [fine_offd.ncols];
    {// 非对角部分
        // 每个进程告诉细网格上的邻居，对方在细网格上的halo区所对应的（由我负责的）列号在粗网格中的全局列号
        static_assert(sizeof(idx_t) <= sizeof(calc_t));
        
        for (idx_t s = 0; s < fine_mat.num_sends; s++) {
            const idx_t offset = fine_mat.send_cnts[s],
                        num = fine_mat.send_cnts[s + 1] - offset;
            idx_t * buf = ( idx_t *)(fine_mat.send_buf     + offset);
            const idx_t * glb_rows = fine_mat.send_row_ids + offset;
            #pragma omp parallel for schedule(static)
            for (idx_t k = 0; k < num; k++) {
                const idx_t loc_fid = glb_rows[k] - fine_mat.beg_row,
                            loc_cid = loc_map_f2c[loc_fid];
                buf[k] = loc_cid + coar_mat->beg_row;
            }
            MPI_Isend(buf, num, mpi_idx_type, fine_mat.send_pids[s], 55, fine_mat.comm, fine_mat.send_reqs + s);
        }
        for (idx_t r = 0; r < fine_mat.num_recvs; r++) {
            const idx_t offset = fine_mat.recv_cnts[r],
                        num = fine_mat.recv_cnts[r + 1] - offset;
            idx_t * buf = (idx_t *)(remote_map_f2c + fine_mat.recv_cnts[r]);
            MPI_Irecv(buf, num, mpi_idx_type, fine_mat.recv_pids[r], 55, fine_mat.comm, fine_mat.recv_reqs + r);
        }
        MPI_Waitall(fine_mat.num_sends, fine_mat.send_reqs, MPI_STATUSES_IGNORE);
        MPI_Waitall(fine_mat.num_recvs, fine_mat.recv_reqs, MPI_STATUSES_IGNORE);
        // 此时接收缓冲区中的是细矩阵非对角部分原来的局部列号在粗矩阵中的全局列号
        // 然后构造粗矩阵的col_map_offd (即局部列序号到全局列序号的映射)
        // 但可能原来的2列才对应粗矩阵中的一列，所以需要消重

        const idx_t num = fine_offd.ncols;
        std::unordered_map<idx_t, idx_t> aux_map;
        idx_t * arr = new idx_t [num];
        #pragma omp parallel for schedule(static)
        for (idx_t k = 0; k < num; k++)
            arr[k] = remote_map_f2c[k];
        // 对remote_map排序
        qsort(arr, num, sizeof(idx_t), cmpfunc<idx_t>);
        const idx_t uniq_cnt = uniquify<idx_t, data_t>(num, arr);

        // 构造coar_mat的非对角部分的列映射
        assert(coar_mat->col_map_offd == nullptr);
        coar_mat->col_map_offd = new idx_t [uniq_cnt];
        for (idx_t j = 0; j < uniq_cnt; j++) {
            coar_mat->col_map_offd[j] = arr[j];
            aux_map.emplace(arr[j], j);// 不可多线程化
        }
        // 构造非对角部分粗化所需的f2c，并储存在arr中
        #pragma omp parallel for schedule(static) // STL允许多个线程同时读
        for (idx_t f_loc = 0; f_loc < num; f_loc++) {
            const idx_t c_glb = remote_map_f2c[f_loc];
            assert(aux_map.find(c_glb) != aux_map.end());
            const idx_t c_loc = aux_map[c_glb];
            arr[f_loc] = c_loc;
        }

        aggregate_mat<idx_t, data_t, calc_t>(fine_offd, num_coar, uniq_cnt, rstr_rpt, rstr_cid, arr, coar_mat->offd, alpha);
        delete [] arr; arr = nullptr;
    }
    // 至此粗网格矩阵的除通信元数据外的其它数据都已全部就绪
    
    {// 建立通信元数据
        const idx_t num_sends = fine_mat.num_sends,
                    num_recvs = fine_mat.num_recvs;
        // 在不发生跨进程粗化时，聚合型的AMG应该邻居进程数不变
        {
            coar_mat->num_sends = num_sends;
            coar_mat->send_pids = new int   [num_sends];
            coar_mat->send_cnts = new idx_t [num_sends + 1];
            idx_t * write_buf = new idx_t [fine_mat.send_cnts[num_sends]];// 按最大可能的需求开辟
            // 根据细矩阵的元数据修改得到粗矩阵的
            coar_mat->send_cnts[0] = 0;
            #pragma omp parallel for schedule(static)
            for (idx_t s = 0; s < num_sends; s ++) {
                coar_mat->send_pids[s] = fine_mat.send_pids[s];
                const idx_t offset = fine_mat.send_cnts[s];
                const idx_t * f_glbs = fine_mat.send_row_ids + offset;
                idx_t * dst_buf = write_buf + offset;

                const idx_t f_num = fine_mat.send_cnts[s + 1] - offset;
                for (idx_t fk = 0; fk < f_num; fk++) {
                    const idx_t f_loc = f_glbs[fk] - fine_mat.beg_row,
                                c_loc = loc_map_f2c[f_loc];
                    dst_buf[fk] = c_loc + coar_mat->beg_row;
                }
                qsort(dst_buf, f_num, sizeof(idx_t), cmpfunc<idx_t>);
                const idx_t uniq_cnt = uniquify<idx_t, data_t>(f_num, dst_buf);
                coar_mat->send_cnts[s + 1] = uniq_cnt;
            }
            
            for (idx_t s = 0; s < num_sends; s ++)// 单线程前缀和
                coar_mat->send_cnts[s + 1] += coar_mat->send_cnts[s];
            coar_mat->send_row_ids = new idx_t [coar_mat->send_cnts[num_sends]];
            
            #pragma omp parallel for schedule(static)
            for (idx_t s = 0; s < num_sends; s ++) {
                const idx_t f_offset = fine_mat.send_cnts[s],
                            c_offset = coar_mat->send_cnts[s];
                const idx_t c_num = coar_mat->send_cnts[s + 1] - c_offset;
                for (idx_t j = 0; j < c_num; j++)
                    coar_mat->send_row_ids[c_offset + j] = write_buf[f_offset + j];
            }
            delete [] write_buf; write_buf = nullptr;
        }
        {
            coar_mat->num_recvs = num_recvs;
            coar_mat->recv_pids = new int   [num_recvs];
            coar_mat->recv_cnts = new idx_t [num_recvs + 1];
            // 根据细矩阵的元数据修改得到粗矩阵的
            coar_mat->recv_cnts[0] = 0;
            assert(fine_offd.ncols == fine_mat.recv_cnts[num_recvs]);
            #pragma omp parallel for schedule(static)
            for (idx_t r = 0; r < num_recvs; r++) {
                coar_mat->recv_pids[r] = fine_mat.recv_pids[r];
                const idx_t offset = fine_mat.recv_cnts[r];
                const idx_t * c_glbs = remote_map_f2c + offset;
                const idx_t f_num = fine_mat.recv_cnts[r + 1] - offset;
                idx_t _poss[f_num];
                for (idx_t j = 0; j < f_num; j++)
                    _poss[j] = c_glbs[j];
                qsort(_poss, f_num, sizeof(idx_t), cmpfunc<idx_t>);
                const idx_t uniq_cnt = uniquify<idx_t, data_t>(f_num, _poss);
                coar_mat->recv_cnts[r + 1] = uniq_cnt;
            }
            for (idx_t r = 0; r < num_recvs; r++)
                coar_mat->recv_cnts[r + 1] += coar_mat->recv_cnts[r];
        }
        coar_mat->send_reqs = new MPI_Request [num_sends];
        coar_mat->recv_reqs = new MPI_Request [num_recvs];
        coar_mat->send_buf = new calc_t [coar_mat->send_cnts[num_sends]];
    }
    delete [] remote_map_f2c; remote_map_f2c = nullptr;

    // 返回的所有粗化信息
    CoarsenWrapper<idx_t, data_t, calc_t> * ret = new CoarsenWrapper<idx_t, data_t, calc_t>;
    ret->loc_map_f2c = loc_map_f2c;
    ret->R_rp = rstr_rpt;
    ret->R_ci = rstr_cid;
    ret->coar_mat = coar_mat;

    return ret;
}

#endif
